package self

import (
	"database/sql"

	"github.com/kataras/iris"

	"github.com/lib/pq"
)

// Defender defend the postgres
type Defender struct{ *PgSelf }

// Watch watch the contex
func (s *PgSelf) Watch(ctx iris.Context) *Defender {
	s.ctx = ctx
	return &Defender{s}
}

// Defend defend the postgreSQL query
func (s *Defender) Defend(handle func(pg *sql.DB) error) bool {
	var rolename string
	if s.claims == nil {
		rolename = "nouser"
	} else {
		rolename = s.claims.Rolename
	}

	// Set rolename
	_, err := s.pg.Exec("SET ROLE " + rolename)
	if err != nil {
		if pgErr, ok := err.(*pq.Error); ok {
			panic("exec set role get error:" + pgErr.Code.Name())
		} else {
			panic("exec set role get unknown error:" + err.Error())
		}
	} else {
		s.logf("current rolename is " + rolename)
	}

	err = handle(s.pg)
	if err != nil {
		if pgErr, ok := err.(*pq.Error); ok {
			s.logf("defending get error:%s with rolename %s", err.Error(), rolename)
			switch pgErr.Code {
			case "42501":
				s.ctx.StatusCode(iris.StatusForbidden)
				s.ctx.Values().Set("message", "role "+rolename+" is not allow to query data")
				return false
			}
		} else {
			s.logf("defending get other error:%v with rolename %s", err, rolename)
		}
		if err != sql.ErrNoRows {
			s.ctx.StatusCode(iris.StatusInternalServerError)
			s.ctx.Values().Set("message", "when query data get unkown server error!")
		}
		return false
	}

	s.pg.Exec("RESET ROLE")
	return true
}
